import torch
from torch import nn
from transformers import BertTokenizer, BigBirdModel
import model.config as conf
from torch.nn import functional

class BigBird(nn.Module):
    def __init__(self, n_model, cls_nu):
        super(BigBird, self).__init__()
        # tokenizer = AutoTokenizer.from_pretrained("schen/longformer-chinese-base-4096")
        self.lf_model = BigBirdModel.from_pretrained(conf.model_name_or_path)
        # self.dropout = nn.Dropout(p=0.2)
        self.fc = nn.Linear(n_model, 1)  # 应用2分类
        # self.fc = nn.Linear(n_model, cls_nu)

    def forward(self, x, atten_mask):
        result = self.lf_model(x, atten_mask)
        # result = self.dropout(result.prediction_logits[:, 0, :]) # longformer
        # result = self.dropout(result.pooler_output)  # bigbird
        # result = self.dropout(result)  # bigbird
        result = self.fc(result.pooler_output)
        # result = functional.softmax(result, dim=1)
        # result = torch.sigmoid(result)
        return result
